import React, { useEffect, useState } from "react";
import { Space, Tree, Checkbox } from "antd";
import flatten from "@/utils/flatten";
import PropTypes from "prop-types";

const TreeSelectComponent = (props) => {
  const [checkedKeys, setCheckedKeys] = useState([]);
  const [expandedKeys, setExpandedKeys] = useState([]);
  const [strictly, setStrictly] = useState(false);
  const [checkAll, setCheckAll] = useState(false);
  const [expandParent, setExpandParent] = useState(false);

  const treeKeys = props.treeData ? flatten(props.treeData).map((item) => String(item.key)) : [];

  useEffect(() => {
    props.value && setCheckedKeys(props.value);
  }, [props]);

  const onCheck = (values) => {
    const checkKeys = values.checked || values;
    setCheckedKeys(checkKeys);
    props.onChange && props.onChange(checkKeys);
    setCheckAll(checkKeys.length === treeKeys.length);
  };

  const onExpand = (expandedKeys) => {
    if (expandedKeys) {
      setExpandedKeys(expandedKeys);
    } else {
      setExpandedKeys(expandParent ? [] : treeKeys);
      setExpandParent(!expandParent);
    }
  };

  const onCheckAll = () => {
    const values = checkAll ? [] : treeKeys;
    setCheckedKeys(values);
    setCheckAll(!checkAll);
    props.onChange && props.onChange(values);
  };

  return (
    <>
      <Space style={{ padding: "4px 8px", marginBottom: "8px" }}>
        <Checkbox checked={checkAll} onChange={onCheckAll}>
          {checkAll ? "取消全选" : "全选"}
        </Checkbox>
        <Checkbox checked={expandParent} onChange={() => onExpand()}>
          展开/折叠
        </Checkbox>
        <Checkbox checked={strictly} onChange={() => setStrictly(!strictly)}>
          父子联动
        </Checkbox>
      </Space>
      <Tree
        {...props}
        checkable
        blockNode
        height={160}
        checkedKeys={checkedKeys}
        expandedKeys={expandedKeys}
        autoExpandParent={false}
        checkStrictly={!strictly}
        onCheck={onCheck}
        onExpand={onExpand}
      />
    </>
  );
};
TreeSelectComponent.propTypes = {
  treeData: PropTypes.array,
  value: PropTypes.array,
  onChange: PropTypes.func,
};

export default TreeSelectComponent;
